import win32gui
from ctypes import windll, byref, c_ubyte
from ctypes.wintypes import RECT
import keyboard
import numpy as np
import cv2
import os
import time

# 排除缩放干扰
windll.user32.SetProcessDPIAware()

class My_Window:
    def __init__(self, window_title):
        """
        初始化窗口捕获类。
        
        :param window_title: 窗口标题（如 "雷电模拟器"）
        """
        self.handle = win32gui.FindWindow(None, window_title)
        if self.handle == 0:
            
            raise ValueError(f"窗口 '{window_title}' 未找到!")

    def capture(self):
        """
        捕获窗口画面。
        :return: 返回一个 numpy.ndarray 格式的图像（BGR 格式）
        """
        # 获取窗口大小
        r = RECT()
        windll.user32.GetWindowRect(self.handle, byref(r))
        width = r.right - r.left
        height = r.bottom - r.top

        # 开始截图
        dc = windll.user32.GetDC(self.handle)
        cdc = windll.gdi32.CreateCompatibleDC(dc)
        bitmap = windll.gdi32.CreateCompatibleBitmap(dc, width, height)
        windll.gdi32.SelectObject(cdc, bitmap)
        windll.gdi32.BitBlt(cdc, 0, 0, width, height, dc, 0, 0, 0x00CC0020)  # SRCCOPY

        # 获取截图数据
        total_bytes = width * height * 4
        buffer = bytearray(total_bytes)
        byte_array = (c_ubyte * total_bytes).from_buffer(buffer)
        windll.gdi32.GetBitmapBits(bitmap, total_bytes, byte_array)

        # 释放资源
        windll.gdi32.DeleteObject(bitmap)
        windll.gdi32.DeleteObject(cdc)
        windll.user32.ReleaseDC(self.handle, dc)

        frame = np.frombuffer(buffer, dtype=np.uint8).reshape(height, width, 4)
        # 将 BGRA 图像转换为 BGR 格式（YOLO 模型需要 BGR 格式）
        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)
        # 返回 numpy.ndarray 格式的图像
        return frame_bgr
    def get_windonw_xy(self,x,y,w,h):
        """
        获取图片的区域截图,并返回截图结果
        """
        frame = self.capture()
        image_xy = frame[y:y+h,x:x+w]

        return image_xy



    def generate_unique_filename(self, save_path, base_name=""):
        """
        生成不重复的文件名。
        
        :param save_path: 文件保存路径
        :param base_name: 文件名前缀，若为空则使用数字编号
        :return: 不重复的文件完整路径
        """
        existing_files = set(os.listdir(save_path))  # 获取目录下的所有文件
        base_name = base_name or ""  # 如果为空，则用空字符串
        index = 1

        while True:
            # 根据 base_name 构造文件名
            if base_name:
                filename = f"{base_name}{index}.png"
            else:
                filename = f"{index}.png"
            
            # 检查文件是否存在
            if filename not in existing_files:
                return os.path.join(save_path, filename)
            
            index += 1  # 文件名递增  
    def key_capture(self, trigger_key, save_path, base_name=""):
        """
        监听按键,按下指定的触发键后截图并保存,按下ESC退出。
        :param trigger_key: 截图的触发按键（如 'q')
        :param save_path: 保存路径
        :param base_name: 保存的文件名前缀
        """
        # 确保保存路径存在
        os.makedirs(save_path, exist_ok=True)
        print(f"监听按键 '{trigger_key}' 进行截图，按 'ESC' 退出。")

        try:
            while True:
                # 如果按下触发键，进行截图
                if keyboard.is_pressed(trigger_key):
                    # 生成唯一文件名
                    file_path = self.generate_unique_filename(save_path, base_name)                   
                    # 调用 capture 方法截图
                    screenshot = self.capture()  # numpy.ndarray 格式
                    # 保存为 PNG 格式
                    cv2.imwrite(file_path, screenshot)
                    # image.save(file_path)
                    print(f"截图已保存到: {file_path}")                    
                    # 防止多次触发，等待键释放
                    while keyboard.is_pressed(trigger_key):
                        time.sleep(0.1)
                # 如果按下 ESC 键，退出
                if keyboard.is_pressed("esc"):
                    print("已退出截图模式。")
                    break
        except KeyboardInterrupt:
            print("程序中断退出。")
    def get_window_rect(self):
            """
            获取窗口的坐标和大小。
                """
            left, top, right, bottom = win32gui.GetWindowRect(self.handle)
            width = right - left
            height = bottom - top
            return left, top, width, height
            #return width, height


if __name__ == "__main__":
    window = My_Window("雷电模拟器")
    x = window.get_windonw_xy(189,258,362,123)
    cv2.imwrite("test.png",x)
    print("截图完成")

